import math
from math import cos, sin
import numpy as np
import rospy
from tf2_geometry_msgs import do_transform_pose
import tf2_ros
from geometry_msgs.msg import Pose, PoseStamped

def rot_X(theta):
    return np.array([[1,0,0],[0,cos(theta),-sin(theta)],[0,sin(theta),cos(theta)]],dtype=np.float32)
def rot_Y(theta):
    return np.array([[cos(theta),0,sin(theta)],[0,1,0],[-sin(theta),0,cos(theta)]],dtype=np.float32)
def rot_Z(theta):
    return np.array([[cos(theta),-sin(theta),0],[sin(theta),cos(theta),0],[0,0,1]],dtype=np.float32)

def get_pose_list_from_pose_stamped(pose_stamped):
    position = [pose_stamped.pose.position.x, pose_stamped.pose.position.y, pose_stamped.pose.position.z]
    orientation = [pose_stamped.pose.orientation.x, pose_stamped.pose.orientation.y,
                   pose_stamped.pose.orientation.z, pose_stamped.pose.orientation.w]
    return position + orientation

def R2q(R):
    w = R[0,0]+R[1,1]+R[2,2]+1
    x = R[0,0]-R[1,1]-R[2,2]+1
    y = -R[0,0]+R[1,1]-R[2,2]+1
    z = -R[0,0]-R[1,1]+R[2,2]+1
    q = np.array([w,x,y,z])
    index = np.argmax(q)
    # assert q[index]>0, "max(q) > 0"
    if q[index]<=0:
        return np.array([0,0,0,1],dtype=np.float32)
    q[index] = math.sqrt(q[index]) / 2
    if index==0:
        q0 = q[index]
        q1 = (R[2,1]-R[1,2]) / (4*q0)
        q2 = (R[0,2]-R[2,0]) / (4*q0)
        q3 = (R[1,0]-R[0,1]) / (4*q0)
    elif index==1:
        q1 = q[index]
        q0 = (R[2,1]-R[1,2]) / (4*q1)
        q2 = (R[0,1]+R[1,0]) / (4*q1)
        q3 = (R[2,0]+R[0,2]) / (4*q1)
    elif index==2:
        q2 = q[index]
        q0 = (R[0,2]-R[2,0]) / (4*q2)
        q1 = (R[0,1]+R[1,0]) / (4*q2)
        q3 = (R[1,2]+R[2,1]) / (4*q2)
    elif index==3:
        q3 = q[index]
        q0 = (R[1,0]-R[0,1]) / (4*q3)
        q1 = (R[2,0]+R[0,2]) / (4*q3)
        q2 = (R[1,2]+R[2,1]) / (4*q3)
    else:
        raise ValueError('index error:' +  str(index))
    # print(np.linalg.norm(np.array([q1,q2,q3,q0],dtype=np.float32)))
    return np.array([q1,q2,q3,q0],dtype=np.float32)



def transform_posestamped_to_tgt(pose, tgt, tfBuffer):
    '''
    args:
        pose: PoseStamped, pose that needs to be transformed
        tgt: str, target frame name
        tfBuffer: tf2_ros.Buffer()
    return:
        transformed_pose: PoseStamped
    '''
    transform_ok = False
    while not transform_ok and not rospy.is_shutdown():
        try:
            transform = tfBuffer.lookup_transform(tgt, pose.header.frame_id, rospy.Time(0))
            transformed_pose = do_transform_pose(pose, transform)
            transform_ok = True
        except tf2_ros.ExtrapolationException as e:
            rospy.logwarn( "Exception on transforming point... trying again \n(" +str(e) + ")")
            rospy.sleep(0.01)
            pose.header.stamp = tfBuffer.get_latest_common_time(tgt, pose.header.frame_id)

    return transformed_pose

def transform_poselist_to_tgt(pose_list, tgt, tfBuffer):
    '''
    args:
        pose: list of PoseStamped, pose that needs to be transformed
        tgt: str, target frame name
        tfBuffer: tf2_ros.Buffer()
    return:
        transformed_pose: PoseStamped
    '''
    assert isinstance(pose_list, list)
    transform_ok = False
    while not transform_ok and not rospy.is_shutdown():
        try:
            transformed_pose_list = []
            for pose in pose_list:
                transform = tfBuffer.lookup_transform(tgt, pose.header.frame_id, rospy.Time(0))
                transformed_pose = do_transform_pose(pose, transform)
                transformed_pose_list.append(transformed_pose)
            transform_ok = True
        except tf2_ros.ExtrapolationException as e:
            rospy.logwarn( "Exception on transforming point... trying again \n(" +str(e) + ")")
            rospy.sleep(0.01)
            pose.header.stamp = tfBuffer.get_latest_common_time(tgt, pose.header.frame_id)

    return transformed_pose_list


def get_pos_and_orient_list(Rt):
    '''
    args:
        Rt: np.ndarray(4x4)
    return:
        pos: list
        orient: list
    '''
    
    position = Rt[:3,3].tolist()
    orientation = R2q(Rt[:3,:3]).tolist()

    return position, orientation